{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: torch==2.2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (2.2.0)\n",
      "Requirement already satisfied: filelock in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (3.13.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (4.9.0)\n",
      "Requirement already satisfied: sympy in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (1.12)\n",
      "Requirement already satisfied: networkx in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (2.8.5)\n",
      "Requirement already satisfied: jinja2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (3.1.3)\n",
      "Requirement already satisfied: fsspec in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (2024.2.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from jinja2->torch==2.2) (2.1.5)\n",
      "Requirement already satisfied: mpmath>=0.19 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from sympy->torch==2.2) (1.3.0)\n",
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.0 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: torchvision==0.17 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (0.17.0)\n",
      "Requirement already satisfied: numpy in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torchvision==0.17) (1.26.4)\n",
      "Requirement already satisfied: requests in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torchvision==0.17) (2.31.0)\n",
      "Requirement already satisfied: torch==2.2.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torchvision==0.17) (2.2.0)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torchvision==0.17) (9.3.0)\n",
      "Requirement already satisfied: filelock in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchvision==0.17) (3.13.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchvision==0.17) (4.9.0)\n",
      "Requirement already satisfied: sympy in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchvision==0.17) (1.12)\n",
      "Requirement already satisfied: networkx in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchvision==0.17) (2.8.5)\n",
      "Requirement already satisfied: jinja2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchvision==0.17) (3.1.3)\n",
      "Requirement already satisfied: fsspec in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchvision==0.17) (2024.2.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests->torchvision==0.17) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests->torchvision==0.17) (3.6)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests->torchvision==0.17) (2.2.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests->torchvision==0.17) (2024.2.2)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from jinja2->torch==2.2.0->torchvision==0.17) (2.1.5)\n",
      "Requirement already satisfied: mpmath>=0.19 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from sympy->torch==2.2.0->torchvision==0.17) (1.3.0)\n",
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.0 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: matplotlib==3.5.2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (3.5.2)\n",
      "Requirement already satisfied: cycler>=0.10 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from matplotlib==3.5.2) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from matplotlib==3.5.2) (4.49.0)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from matplotlib==3.5.2) (1.4.5)\n",
      "Requirement already satisfied: numpy>=1.17 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from matplotlib==3.5.2) (1.26.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from matplotlib==3.5.2) (23.2)\n",
      "Requirement already satisfied: pillow>=6.2.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from matplotlib==3.5.2) (9.3.0)\n",
      "Requirement already satisfied: pyparsing>=2.2.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from matplotlib==3.5.2) (3.1.1)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from matplotlib==3.5.2) (2.8.2)\n",
      "Requirement already satisfied: six>=1.5 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from python-dateutil>=2.7->matplotlib==3.5.2) (1.16.0)\n",
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.0 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: scikit-image==0.19.3 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (0.19.3)\n",
      "Requirement already satisfied: numpy>=1.17.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from scikit-image==0.19.3) (1.26.4)\n",
      "Requirement already satisfied: scipy>=1.4.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from scikit-image==0.19.3) (1.12.0)\n",
      "Requirement already satisfied: networkx>=2.2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from scikit-image==0.19.3) (2.8.5)\n",
      "Requirement already satisfied: pillow!=7.1.0,!=7.1.1,!=8.3.0,>=6.1.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from scikit-image==0.19.3) (9.3.0)\n",
      "Requirement already satisfied: imageio>=2.4.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from scikit-image==0.19.3) (2.34.0)\n",
      "Requirement already satisfied: tifffile>=2019.7.26 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from scikit-image==0.19.3) (2024.2.12)\n",
      "Requirement already satisfied: PyWavelets>=1.1.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from scikit-image==0.19.3) (1.5.0)\n",
      "Requirement already satisfied: packaging>=20.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from scikit-image==0.19.3) (23.2)\n",
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.0 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: pytorch-lightning==1.7.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (1.7.0)\n",
      "Requirement already satisfied: numpy>=1.17.2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from pytorch-lightning==1.7.0) (1.26.4)\n",
      "Requirement already satisfied: torch>=1.9.* in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from pytorch-lightning==1.7.0) (2.2.0)\n",
      "Requirement already satisfied: tqdm>=4.57.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from pytorch-lightning==1.7.0) (4.66.2)\n",
      "Requirement already satisfied: PyYAML>=5.4 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from pytorch-lightning==1.7.0) (6.0.1)\n",
      "Requirement already satisfied: fsspec!=2021.06.0,>=2021.05.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.0) (2024.2.0)\n",
      "Requirement already satisfied: tensorboard>=2.9.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from pytorch-lightning==1.7.0) (2.15.2)\n",
      "Requirement already satisfied: torchmetrics>=0.7.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from pytorch-lightning==1.7.0) (0.11.4)\n",
      "Requirement already satisfied: pyDeprecate>=0.3.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from pytorch-lightning==1.7.0) (0.3.2)\n",
      "Requirement already satisfied: packaging>=17.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from pytorch-lightning==1.7.0) (23.2)\n",
      "Requirement already satisfied: typing-extensions>=4.0.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from pytorch-lightning==1.7.0) (4.9.0)\n",
      "Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.0) (3.9.3)\n",
      "Requirement already satisfied: absl-py>=0.4 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.0) (2.1.0)\n",
      "Requirement already satisfied: grpcio>=1.48.2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.0) (1.60.1)\n",
      "Requirement already satisfied: google-auth<3,>=1.6.3 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.0) (2.28.0)\n",
      "Requirement already satisfied: google-auth-oauthlib<2,>=0.5 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.0) (1.2.0)\n",
      "Requirement already satisfied: markdown>=2.6.8 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.0) (3.5.2)\n",
      "Requirement already satisfied: protobuf!=4.24.0,>=3.19.6 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.0) (4.25.3)\n",
      "Requirement already satisfied: requests<3,>=2.21.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.0) (2.31.0)\n",
      "Requirement already satisfied: setuptools>=41.0.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.0) (68.2.2)\n",
      "Requirement already satisfied: six>1.9 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.0) (1.16.0)\n",
      "Requirement already satisfied: tensorboard-data-server<0.8.0,>=0.7.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.0) (0.7.2)\n",
      "Requirement already satisfied: werkzeug>=1.0.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from tensorboard>=2.9.1->pytorch-lightning==1.7.0) (3.0.1)\n",
      "Requirement already satisfied: filelock in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch>=1.9.*->pytorch-lightning==1.7.0) (3.13.1)\n",
      "Requirement already satisfied: sympy in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch>=1.9.*->pytorch-lightning==1.7.0) (1.12)\n",
      "Requirement already satisfied: networkx in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch>=1.9.*->pytorch-lightning==1.7.0) (2.8.5)\n",
      "Requirement already satisfied: jinja2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch>=1.9.*->pytorch-lightning==1.7.0) (3.1.3)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.0) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.0) (23.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.0) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.0) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.0) (1.9.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning==1.7.0) (4.0.3)\n",
      "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (5.3.2)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (0.3.0)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (4.9)\n",
      "Requirement already satisfied: requests-oauthlib>=0.7.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from google-auth-oauthlib<2,>=0.5->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (1.3.1)\n",
      "Requirement already satisfied: importlib-metadata>=4.4 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from markdown>=2.6.8->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (7.0.1)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests<3,>=2.21.0->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests<3,>=2.21.0->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (3.6)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests<3,>=2.21.0->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (2.2.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests<3,>=2.21.0->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (2024.2.2)\n",
      "Requirement already satisfied: MarkupSafe>=2.1.1 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from werkzeug>=1.0.1->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (2.1.5)\n",
      "Requirement already satisfied: mpmath>=0.19 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from sympy->torch>=1.9.*->pytorch-lightning==1.7.0) (1.3.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (3.17.0)\n",
      "Requirement already satisfied: pyasn1<0.6.0,>=0.4.6 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (0.5.1)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<2,>=0.5->tensorboard>=2.9.1->pytorch-lightning==1.7.0) (3.2.2)\n",
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.0 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install torch==2.2\n",
    "!pip install torchvision==0.17\n",
    "!pip install matplotlib==3.5.2\n",
    "!pip install scikit-image==0.19.3\n",
    "!pip install pytorch-lightning==1.7.0\n",
    "# !pip install torchmetrics==0.11.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.datasets import MNIST\n",
    "from torchvision import transforms\n",
    "import pytorch_lightning as pl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvNet(pl.LightningModule):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(ConvNet, self).__init__()\n",
    "        self.cn1 = nn.Conv2d(1, 16, 3, 1)\n",
    "        self.cn2 = nn.Conv2d(16, 32, 3, 1)\n",
    "        self.dp1 = nn.Dropout2d(0.10)\n",
    "        self.dp2 = nn.Dropout2d(0.25)\n",
    "        self.fc1 = nn.Linear(4608, 64) # 4608 is basically 12 X 12 X 32\n",
    "        self.fc2 = nn.Linear(64, 10)\n",
    " \n",
    "    def forward(self, x):\n",
    "        x = self.cn1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.cn2(x)\n",
    "        x = F.relu(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = self.dp1(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.dp2(x)\n",
    "        x = self.fc2(x)\n",
    "        op = F.log_softmax(x, dim=1)\n",
    "        return op\n",
    "\n",
    "    def training_step(self, batch, batch_num):\n",
    "        train_x, train_y = batch\n",
    "        y_pred = self(train_x)\n",
    "        training_loss = F.cross_entropy(y_pred, train_y)\n",
    "        # optional\n",
    "        self.log('train_loss', training_loss, on_epoch=True, prog_bar=True)\n",
    "        return training_loss\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        # optional\n",
    "        avg_loss = torch.stack(outputs).mean()\n",
    "        return avg_loss\n",
    "\n",
    "    def test_step(self, batch, batch_num):\n",
    "        # optional\n",
    "        test_x, test_y = batch\n",
    "        y_pred = self(test_x)\n",
    "        test_loss = F.cross_entropy(y_pred, test_y)\n",
    "        # optional\n",
    "        self.log('test_loss', test_loss, on_step=True, on_epoch=True, prog_bar=True)\n",
    "        return test_loss\n",
    "\n",
    "    def test_epoch_end(self, outputs):\n",
    "        # optional\n",
    "        avg_loss = torch.stack(outputs).mean()\n",
    "        return avg_loss\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adadelta(self.parameters(), lr=0.5)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(MNIST(os.getcwd(), train=True, download=True, \n",
    "                                transform=transforms.Compose([transforms.ToTensor(),\n",
    "                                                              transforms.Normalize((0.1302,), (0.3069,))])), \n",
    "                                batch_size=32, num_workers=4)\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        # optional\n",
    "        return DataLoader(MNIST(os.getcwd(), train=True, download=True, \n",
    "                                transform=transforms.Compose([transforms.ToTensor(),\n",
    "                                                              transforms.Normalize((0.1302,), (0.3069,))])), \n",
    "                                batch_size=32, num_workers=4)\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        # optional\n",
    "        return DataLoader(MNIST(os.getcwd(), train=False, download=True, \n",
    "                                transform=transforms.Compose([transforms.ToTensor(),\n",
    "                                                              transforms.Normalize((0.1302,), (0.3069,))])), \n",
    "                                batch_size=32, num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1788: UserWarning: MPS available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='mps', devices=1)`.\n",
      "  rank_zero_warn(\n",
      "/Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages/pytorch_lightning/trainer/configuration_validator.py:115: UserWarning: You passed in a `val_dataloader` but have no `validation_step`. Skipping val loop.\n",
      "  rank_zero_warn(\"You passed in a `val_dataloader` but have no `validation_step`. Skipping val loop.\")\n",
      "\n",
      "  | Name | Type      | Params\n",
      "-----------------------------------\n",
      "0 | cn1  | Conv2d    | 160   \n",
      "1 | cn2  | Conv2d    | 4.6 K \n",
      "2 | dp1  | Dropout2d | 0     \n",
      "3 | dp2  | Dropout2d | 0     \n",
      "4 | fc1  | Linear    | 294 K \n",
      "5 | fc2  | Linear    | 650   \n",
      "-----------------------------------\n",
      "300 K     Trainable params\n",
      "0         Non-trainable params\n",
      "300 K     Total params\n",
      "1.202     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b1eeee117e424c98b7c44541c9fe52c0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages/torch/nn/functional.py:1347: UserWarning: dropout2d: Received a 2-D input to dropout2d, which is deprecated and will result in an error in a future release. To retain the behavior and silence this warning, please use dropout instead. Note that dropout2d exists to provide channel-wise dropout on inputs with 2 spatial dimensions, a channel dimension, and an optional batch dimension (i.e. 3D or 4D inputs).\n",
      "  warnings.warn(warn_msg)\n",
      "`Trainer.fit` stopped: `max_epochs=10` reached.\n"
     ]
    }
   ],
   "source": [
    "model = ConvNet()\n",
    "\n",
    "trainer = pl.Trainer(max_epochs=10)    \n",
    "trainer.fit(model)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1385: UserWarning: `.test(ckpt_path=None)` was called without a model. The best model of the previous `fit` call will be used. You can pass `.test(ckpt_path='best')` to use the best model or `.test(ckpt_path='last')` to use the last model. If you pass a value, this warning will be silenced.\n",
      "  rank_zero_warn(\n",
      "Restoring states from the checkpoint path at /Users/Ashish.Jha/code/personal/MasteringPyTorchV2/Chapter15/lightning_logs/version_14/checkpoints/epoch=9-step=18750.ckpt\n",
      "Loaded model weights from checkpoint at /Users/Ashish.Jha/code/personal/MasteringPyTorchV2/Chapter15/lightning_logs/version_14/checkpoints/epoch=9-step=18750.ckpt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0b7cfe45d614ce3926f0c20f338ffaa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "     test_loss_epoch        0.03324906900525093\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss_epoch': 0.03324906900525093}]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "      <iframe id=\"tensorboard-frame-d41c6264dc1a7fcc\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
       "      </iframe>\n",
       "      <script>\n",
       "        (function() {\n",
       "          const frame = document.getElementById(\"tensorboard-frame-d41c6264dc1a7fcc\");\n",
       "          const url = new URL(\"/\", window.location);\n",
       "          const port = 6006;\n",
       "          if (port) {\n",
       "            url.port = port;\n",
       "          }\n",
       "          frame.src = url;\n",
       "        })();\n",
       "      </script>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Start tensorboard.\n",
    "%reload_ext tensorboard\n",
    "%tensorboard --logdir lightning_logs/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
